module D =
struct

  module Vec =
  struct

    let nrm1 x =
      Lacaml.D.asum x;;
    
    let partition p x =
      let n = Bigarray.Array1.dim x in
      let rec part_indices i yes no =
        if ( i > n ) then
          (List.rev yes, List.rev no)
        else
          if ( p x.{i} ) then 
            part_indices (i+1) (i::yes) no
          else 
            part_indices (i+1) yes (i::no) in
      let (yes,no) = part_indices 1 [] [] in
      let n1  = List.length yes and n2 = List.length no in
      assert (n = n1 + n2);
      let n2 = (n - n1) in
      let x1 = Lacaml.D.Vec.create n1 in
      let x2 = Lacaml.D.Vec.create n2 in
      begin
        let acc = ref 0 in
          List.iter ( fun idx -> acc := !acc + 1; x1.{!acc} <- x.{idx} ) yes;
        let acc = ref 0 in
          List.iter ( fun idx -> acc := !acc + 1; x2.{!acc} <- x.{idx} ) no;
        (x1,x2);
      end;;

    let for_all p x =
      Lacaml.D.Vec.fold (fun b t -> (p t) && b) true x;;

    let randn ?rnd_state ?(mu = 0.0) ?(sigma = 1.0) n =
      let vec = Lacaml.D.Vec.create n in
      let state =
        match rnd_state with
        | None -> Random.get_state ()
        | Some state -> state in
      for row = 1 to n do
        vec.{row} <- MLlab.Random.randn ~rnd_state:state mu sigma
      done;
      if rnd_state = None then Random.set_state state;
      vec

    (* returns the first index i of the array such that p x.{i} is true *)
    let findi p x =
      let n = Lacaml.D.Vec.dim x in
      let idx  = ref 0 in
      let cond = ref false in
      while ( not !cond && !idx < n ) do
        idx  := !idx + 1;
        cond := p x.{!idx};
      done;
      if !cond then
        !idx
      else
        raise Not_found;;

  end

  module Mat =
  struct
    let randn ?rnd_state ?(mu = 0.0) ?(sigma = 1.0) m n =
      let mat = Lacaml.D.Mat.create m n in
      let state =
        match rnd_state with
        | None -> Random.get_state ()
        | Some state -> state in
      for row = 1 to m do
        for col = 1 to n do
          mat.{row, col} <- MLlab.Random.randn ~rnd_state:state mu sigma
        done
      done;
      if rnd_state = None then Random.set_state state;
      mat
  end

end
